import pymysql
from datetime import datetime

class DBUtils:
    def __init__(self):
        self.connection = None
        
    def connect(self):
        """连接数据库"""
        try:
            self.connection = pymysql.connect(
                host='localhost',
                user='root',
                password='123456',
                database='military',
                charset='utf8mb4'
            )
            return True
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            return False
            
    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            
    def insert_detection(self, detection_data):
        """插入检测数据"""
        if not self.connection:
            if not self.connect():
                return False
                
        try:
            with self.connection.cursor() as cursor:
                sql = """
                INSERT INTO detections (
                    media_id, file_det_name, target_class, pre_cision,
                    detect_time, locate_time, x_center, y_center, width, height
                ) VALUES (
                    %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
                )
                """
                
                cursor.execute(sql, (
                    detection_data['media_id'],
                    detection_data['file_det_name'],
                    detection_data['target_class'],
                    detection_data['pre_cision'],
                    detection_data['detect_time'],
                    detection_data['locate_time'],
                    detection_data['x_center'],
                    detection_data['y_center'],
                    detection_data['width'],
                    detection_data['height']
                ))
                
            self.connection.commit()
            return True
            
        except Exception as e:
            print(f"插入数据失败: {str(e)}")
            self.connection.rollback()
            return False 